package main

import (
	"context"
	"fmt"
	"log"
	"sync"
	"time"

	clientv3 "go.etcd.io/etcd/client/v3"
	"go.etcd.io/etcd/client/v3/concurrency"
)

type DistributedRWLock struct {
	client   *clientv3.Client
	session  *concurrency.Session
	writeKey string
	writeMux *concurrency.Mutex //新增字段保存锁实例
	readKeys sync.Map           // map[string]struct{} to track active read keys (conceptual)
}

func NewDistributedRWLock(endpoints []string, lockPrefix string) (*DistributedRWLock, error) {
	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   endpoints,
		DialTimeout: 5 * time.Second,
	})
	if err != nil {
		return nil, err
	}

	session, err := concurrency.NewSession(cli, concurrency.WithTTL(5))
	if err != nil {
		cli.Close()
		return nil, err
	}

	return &DistributedRWLock{
		client:   cli,
		session:  session,
		writeKey: lockPrefix + "/write",
	}, nil
}

func (rwl *DistributedRWLock) RLock(ctx context.Context, key string) error {
	// In a real implementation, you might want to use etcd to manage read locks per key
	// Here, we're just simulating with a local sync.Map for demonstration
	// To implement a true distributed read lock, consider using leases and etcd keys to track readers

	// For simplicity, we'll just store the key in a sync.Map (not truly distributed)
	_, loaded := rwl.readKeys.LoadOrStore(key, struct{}{})
	if loaded {
		// Key already exists, meaning another reader has the lock (conceptual)
		// In a real distributed scenario, you'd need to check etcd for actual read locks
		return nil // For this demo, we allow multiple "readers" conceptually
	}

	fmt.Printf("Read lock acquired on key: %s\n", key)
	return nil
}

func (rwl *DistributedRWLock) RUnlock(ctx context.Context, key string) {
	rwl.readKeys.Delete(key)
	fmt.Printf("Read lock released on key: %s\n", key)
}

// WLock 在 WLock 方法中保存实例
func (rwl *DistributedRWLock) WLock(ctx context.Context) error {
	rwl.writeMux = concurrency.NewMutex(rwl.session, rwl.writeKey)
	if err := rwl.writeMux.Lock(ctx); err != nil {
		return err
	}
	fmt.Println("Write lock acquired")
	return nil
}

// WUnlock 在 WUnlock 中使用已保存的实例
func (rwl *DistributedRWLock) WUnlock(ctx context.Context) error {
	if err := rwl.writeMux.Unlock(ctx); err != nil {
		return err
	}
	fmt.Println("Write lock released")
	return nil
}

func (rwl *DistributedRWLock) Close() {
	rwl.session.Close()
	rwl.client.Close()
}

func main() {
	endpoints := []string{"localhost:2379"}
	lockPrefix := "/my-distributed-rwlock"

	rwLock, err := NewDistributedRWLock(endpoints, lockPrefix)
	if err != nil {
		log.Fatalf("failed to create distributed rwlock: %v", err)
	}
	defer rwLock.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()

	// Write lock example
	err = rwLock.WLock(ctx)
	if err != nil {
		log.Fatalf("failed to acquire write lock: %v", err)
	}
	// Simulate some work with the write lock
	time.Sleep(5 * time.Second)
	err = rwLock.WUnlock(ctx)
	if err != nil {
		log.Fatalf("failed to release write lock: %v", err)
	}

	// Read lock example
	err = rwLock.RLock(ctx, "some-key")
	if err != nil {
		log.Fatalf("failed to acquire read lock: %v", err)
	}
	// Simulate some work with the read lock
	time.Sleep(3 * time.Second)
	rwLock.RUnlock(ctx, "some-key")
}
